import torch
import torchvision.datasets as datasets
import torchvision.transforms as transforms
import torch.utils.data as data_utils
import cv2

# data
train_data = datasets.MNIST(root="mnist", train=True, transform=transforms.ToTensor(), download=True)
test_data = datasets.MNIST(root="mnist", train=False, transform=transforms.ToTensor(), download=False)

# batchsize
train_loader = data_utils.DataLoader(dataset=train_data, batch_size=64, shuffle=True)
test_loader = data_utils.DataLoader(dataset=test_data, batch_size=64, shuffle=True)


# net
class CNN(torch.nn.Module):

    def __init__(self):
        super(CNN, self).__init__()
        self.conv = torch.nn.Sequential(
            torch.nn.Conv2d(1, 32, kernel_size=5, padding=2),
            torch.nn.BatchNorm2d(32),
            torch.nn.ReLU(),
            torch.nn.MaxPool2d(2)
        )

        self.fc = torch.nn.Linear(14 * 14 * 32, 10)

    def forward(self, x):
        out = self.conv(x)
        out = out.view(out.size()[0], -1)
        out = self.fc(out)
        return out


cnn = torch.load("model/mnist_model.pkl")

# loss
loss_func = torch.nn.CrossEntropyLoss()

# test
loss_test = 0
accuracy = 0
for j, (images_t, labels_t) in enumerate(test_loader):
    images = images_t.cpu()
    labels = labels_t.cpu()
    outputs = cnn(images)
    loss_test += loss_func(outputs, labels)
    _, pred = outputs.max(1)
    accuracy += (pred == labels).sum().item()

    images = images.numpy()
    labels = labels.numpy()
    pred = pred.numpy()

    for idx in range(images.shape[0]):
        im_data = images[idx]
        im_label = labels[idx]
        im_pred = pred[idx]
        im_data = im_data.transpose(1, 2, 0)

        print("label", im_label)
        print("pred", im_pred)
        cv2.imshow("imdata", im_data)
        cv2.waitKey(0)


accuracy = accuracy / len(test_data)
print(accuracy)
